#include "homomorphic_mod.h"
#include "cosine_approx.h"

/*
 * @brief 计算 sin(2πx) 的值，输入是一个复数类型的值。
 *
 * @param x 输入的复数值。
 * @return 返回计算结果，类型为 std::complex<double>。
 */
std::complex<double> sin2pi(const std::complex<double> &x)
{
    return std::sin(6.283185307179586 * x);
}

/*
 * @brief 计算 cos(2πx) 的值，输入是一个复数类型的值。
 *
 * @param x 输入的复数值。
 * @return 返回计算结果，类型为 std::complex<double>。
 */
std::complex<double> cos2pi(const std::complex<double> &x)
{
    return std::cos(6.283185307179586 * x);
}

/*
 * @brief 计算 QDiff 的值。
 *
 * 该函数计算 EvalModLiteral 对象的 QDiff 值，QDiff 是 Q 和其最接近的 2 的幂次方之间的比值。
 *
 * @return 返回计算得到的 QDiff 值。
 */
double EvalModLiteral::QDiff() const
{
    return static_cast<double>(Q) / std::exp2(std::round(std::log2(static_cast<double>(Q))));
}

/*
 * @brief 计算 `EvalModLiteral` 对象的深度。
 *
 * 深度的计算依据以下几个因素：
 * - 如果 `Sine` 是 `SineType::Cos1`，那么方法要求至少有 `2*K-1` 的度数，因此深度计算需要考虑 `SineDeg` 和 `2*K-1` 之间的较大值。
 * - `SineDeg`：正弦多项式的度数，用于计算多项式的复杂度。
 * - `DoubleAngle`：表示是否使用双角公式的标志，如果为 `1`，则增加深度。
 * - `ArcSineDeg`：反正弦多项式的度数，影响深度的计算。
 *
 * 该方法通过计算这些参数的对数（以 2 为底），并将其取整得到深度值。最终深度值是所有部分的总和。
 *
 * @return 返回计算得到的深度值。
 */
int EvalModLiteral::Depth() const
{
    int depth = 0;

    if (Sine == SineType::Cos1) { // this method requires a minimum degree of 2*K-1.
        depth += static_cast<int>(std::ceil(std::log2(static_cast<double>(std::max(SineDeg, 2 * K - 1) + 1))));
    } else {
        depth += static_cast<int>(std::ceil(std::log2(static_cast<double>(SineDeg + 1))));
    }

    depth += DoubleAngle;
    depth += static_cast<int>(std::ceil(std::log2(static_cast<double>(ArcSineDeg + 1))));
    return depth;
}

/*
 * @brief 从 EvalModLiteral 对象创建一个新的 EvalModPoly 对象。
 *
 * 这个函数根据 EvalModLiteral 对象中的多个参数进行处理，构造新的多项式（包括近似的正弦、多项式等）。
 *
 * @param evm EvalModLiteral 对象。
 * @return 返回一个新的 EvalModPoly 对象。
 */
EvalModLiteral::EvalModLiteral(uint64_t q, int levelStart, long scalingFactor, enum SineType SineType,
    double messageRatio, int k, int sineDeg, int doubleAngle, int arcSineDeg)
    : Q(q),
      LevelStart(levelStart),
      ScalingFactor(scalingFactor),
      Sine(SineType),
      MessageRatio(messageRatio),
      K(k),
      SineDeg(sineDeg),
      DoubleAngle(doubleAngle),
      ArcSineDeg(arcSineDeg)
{}

/*
 * @brief 计算多项式的深度。
 *
 * 该函数返回多项式的深度，深度是多项式系数数量的对数（以 2 为底）的上限。
 *
 * @return 返回多项式的深度。
 */
int Polynomial::Depth() const
{
    return static_cast<int>(std::ceil(std::log2(static_cast<double>(Coeffs.size()))));
}

/*
 * @brief 计算多项式的度数。
 *
 * 该函数返回多项式的度数，度数是多项式系数数量减一。
 *
 * @return 返回多项式的度数。
 */
int Polynomial::Degree() const
{
    return Coeffs.size() - 1;
}

/*
 * @brief 多项式的默认构造函数。
 *
 * 初始化多项式的最大度数、引导标志、系数 A 和 B 以及多项式基。
 */
Polynomial::Polynomial() : MaxDeg(0), Lead(true), A(1), B(-1), Basis(PolynomialBasis::ChebyshevBasis) {}

/*
 * @brief 多项式的构造函数。
 *
 * 使用给定的最大度数、系数、引导标志、系数 A 和 B 以及多项式基来初始化多项式。
 *
 * @param maxDeg 最大度数。
 * @param coeffs 多项式系数。
 * @param lead 引导标志。
 * @param a 系数 A。
 * @param b 系数 B。
 * @param basis 多项式基。
 */
Polynomial::Polynomial(int maxDeg, const std::vector<std::complex<double>> &coeffs, bool lead, double a, double b,
    PolynomialBasis basis)
    : MaxDeg(maxDeg), Coeffs(coeffs), Lead(lead), A(a), B(b), Basis(basis)
{}

Polynomial::Polynomial(const vector<std::complex<double>> &coeffs, bool lead, double a, double b)
    : MaxDeg(int(coeffs.size()) - 1), Coeffs(coeffs), Lead(lead), A(a), B(b), Basis(PolynomialBasis::StandardBasis)
{}

/*
 * @brief 多项式的构造函数。
 *
 * 使用给定的系数、引导标志、系数 A 和 B 来初始化多项式。
 *
 * @param coeffs 多项式系数。
 * @param lead 引导标志。
 * @param a 系数 A。
 * @param b 系数 B。
 */
Polynomial::Polynomial(const std::vector<std::complex<double>> &coeffs)
    : MaxDeg(int(coeffs.size()) - 1), Coeffs(coeffs), Lead(true), A(0.0), B(0.0), Basis(PolynomialBasis::StandardBasis)
{}

/*
 * @brief 创建一个新的多项式对象。
 *
 * 该函数复制给定的系数并返回一个新的多项式对象。
 *
 * @param coeffs 多项式系数。
 * @return 返回一个新的多项式对象。
 */
Polynomial *Polynomial::NewPoly(std::vector<std::complex<double>> &coeffs)
{
    std::vector<std::complex<double>> c(coeffs.size());
    std::copy(coeffs.begin(), coeffs.end(), c.begin());
    return new Polynomial(c);
}

/*
 * @brief 计算切比雪夫节点。
 *
 * 该函数计算给定区间 [a, b] 上的 n 个切比雪夫节点。
 *
 * @param n 节点数量。
 * @param a 区间起始值。
 * @param b 区间结束值。
 * @return 返回切比雪夫节点的向量。
 */
std::vector<double> Polynomial::chebyshevNodes(int n, double a, double b)
{
    std::vector<double> u(n);
    double x = 0.5 * (a + b);
    double y = 0.5 * (b - a);
    for (int k = 1; k < n + 1; ++k) {
        u[k - 1] = x + y * std::cos((static_cast<double>(k) - 0.5) * (3.141592653589793 / static_cast<double>(n)));
    }
    return u;
}

/*
 * @brief 计算切比雪夫多项式的系数。
 *
 * 该函数根据给定的节点和函数值计算切比雪夫多项式的系数。
 *
 * @param nodes 切比雪夫节点。
 * @param fi 函数值。
 * @param a 区间起始值。
 * @param b 区间结束值。
 * @return 返回切比雪夫多项式的系数。
 */
std::vector<std::complex<double>> Polynomial::chebyCoeffs(const std::vector<double> &nodes,
    const std::vector<std::complex<double>> &fi, double a, double b)
{
    std::vector<std::complex<double>> coeffs(nodes.size(), std::complex<double>(0.0, 0.0));
    int n = nodes.size();
    for (int i = 0; i < n; ++i) {
        std::complex<double> u = std::complex<double>((2 * nodes[i] - a - b) / (b - a), 0);
        std::complex<double> Tprev = 1;
        std::complex<double> T = u;
        for (int j = 0; j < n; ++j) {
            coeffs[j] += fi[i] * Tprev;
            std::complex<double> Tnext = 2.0 * u * T - Tprev;
            Tprev = T;
            T = Tnext;
        }
    }

    coeffs[0] /= std::complex<double>(static_cast<double>(n), 0);
    for (int i = 1; i < n; ++i) {
        coeffs[i] *= (2.0 / std::complex<double>(static_cast<double>(n), 0));
    }

    return coeffs;
}

/*
 * @brief 近似给定函数的多项式。
 *
 * 该函数使用切比雪夫多项式近似给定区间 [a, b] 上的函数。
 *
 * @param function 要近似的函数。
 * @param a 区间起始值。
 * @param b 区间结束值。
 * @param degree 多项式的度数。
 * @return 返回近似的多项式对象。
 */
Polynomial *Polynomial::Approximate(const std::function<std::complex<double>(std::complex<double>)> &function, double a,
    double b, int degree)
{
    std::vector<double> nodes = chebyshevNodes(degree + 1, a, b);
    std::vector<std::complex<double>> fi(nodes.size());

    for (size_t i = 0; i < nodes.size(); ++i) {
        fi[i] = function(std::complex<double>(nodes[i], 0));
    }

    std::vector<std::complex<double>> coeffs = chebyCoeffs(nodes, fi, a, b);
    Polynomial *pol = NewPoly(coeffs);
    pol->A = a;
    pol->B = b;
    pol->MaxDeg = degree;
    pol->Lead = true;
    pol->Basis = PolynomialBasis::ChebyshevBasis;

    return pol;
}

/*
 * @brief 近似给定函数的多项式。
 *
 * 该函数使用切比雪夫多项式近似给定区间 [a, b] 上的函数。
 *
 * @param function 要近似的函数。
 * @param a 区间起始值。
 * @param b 区间结束值。
 * @param degree 多项式的度数。
 * @return 返回近似的多项式对象。
 */
Polynomial *Polynomial::Approximate(std::complex<double> (*function)(const std::complex<double> &), double a, double b,
    int degree)
{
    std::vector<double> nodes = chebyshevNodes(degree + 1, a, b);
    std::vector<std::complex<double>> fi(nodes.size());

    for (size_t i = 0; i < nodes.size(); ++i) {
        fi[i] = function(std::complex<double>(nodes[i], 0));
    }

    std::vector<std::complex<double>> coeffs = chebyCoeffs(nodes, fi, a, b);
    Polynomial *pol = NewPoly(coeffs);
    pol->A = a;
    pol->B = b;
    pol->MaxDeg = degree;
    pol->Lead = true;
    pol->Basis = PolynomialBasis::ChebyshevBasis;

    return pol;
}

int EvalModPoly::LevelStart() const
{
    return levelStart;
}

double EvalModPoly::ScalingFactor() const
{
    return scalingFactor;
}

double EvalModPoly::ScFac() const
{
    return scFac;
}

double EvalModPoly::MessageRatio() const
{
    return messageRatio;
}

double EvalModPoly::A() const
{
    return sinePoly->A;
}

double EvalModPoly::B() const
{
    return sinePoly->B;
}

double EvalModPoly::K() const
{
    return sinePoly->B * scFac;
}

double EvalModPoly::QDiff() const
{
    return qDiff;
}

/*
 * @brief 从 EvalModLiteral 对象创建一个新的 EvalModPoly 对象。
 *
 * 该函数根据 EvalModLiteral 对象的参数构造一个新的 EvalModPoly 对象。
 *
 * @param evm EvalModLiteral 对象。
 * @return 返回一个新的 EvalModPoly 对象。
 */
EvalModPoly EvalModPoly::NewEvalModPolyFromLiteral(EvalModLiteral evm)
{
    CosineApprox cosine_approx;
    Polynomial *arcSinePoly = nullptr;
    Polynomial *sinePoly = nullptr;
    double sqrt2pi = 0.0;

    double scFac = std::exp2(static_cast<double>(evm.DoubleAngle));
    if (scFac == 0) {
        throw std::runtime_error("Division by zero is not allowed.");
    }
    double qDiff = evm.QDiff();

    if (evm.ArcSineDeg > 0) {
        sqrt2pi = 1.0;

        std::vector<std::complex<double>> coeffs(evm.ArcSineDeg + 1);
        coeffs[1] = std::complex<double>(0.15915494309189535 * qDiff, 0);

        for (int i = 3; i < evm.ArcSineDeg + 1; i += 2) {
            coeffs[i] = coeffs[i - 2] *
                std::complex<double>(static_cast<double>(i * i - 4 * i + 4) / static_cast<double>(i * i - i), 0);
        }

        arcSinePoly = new Polynomial(coeffs);
    } else {
        sqrt2pi = std::pow(0.15915494309189535 * qDiff, 1.0 / scFac);
    }

    if (evm.Sine == SineType::Sin) {
        if (evm.DoubleAngle != 0) {
            throw std::runtime_error("cannot use double angle with SineType == Sin");
        }
        sinePoly = Polynomial::Approximate(sin2pi, -double(evm.K), double(evm.K), evm.SineDeg);
    } else if (evm.Sine == SineType::Cos1) {
        vector<complex<double>> cofs =
            cosine_approx.ApproximateCos(evm.K, evm.SineDeg, evm.MessageRatio, evm.DoubleAngle);
        sinePoly = new Polynomial(cofs, true, -static_cast<double>(evm.K) / scFac, static_cast<double>(evm.K) / scFac);
    } else if (evm.Sine == SineType::Cos2) {
        sinePoly = Polynomial::Approximate(cos2pi, -static_cast<double>(evm.K) / scFac,
            static_cast<double>(evm.K) / scFac, evm.SineDeg);
    } else {
        throw std::runtime_error("invalid SineType");
    }

    for (auto &coeff : sinePoly->Coeffs) {
        coeff *= std::complex<double>(sqrt2pi, 0);
    }

    return { evm.LevelStart, evm.ScalingFactor, evm.Sine, evm.MessageRatio, evm.DoubleAngle, qDiff,
        scFac,          sqrt2pi,           sinePoly, arcSinePoly };
}

/*
 * @brief EvalModPoly 的默认构造函数。
 *
 * 初始化 EvalModPoly 对象的各个成员变量。
 */
EvalModPoly::EvalModPoly()
    : levelStart(0),
      scalingFactor(0),
      sineType(SineType::Cos2),
      messageRatio(0),
      doubleAngle(0),
      qDiff(0),
      scFac(0),
      sqrt2Pi(0),
      sinePoly(nullptr),
      arcSinePoly(nullptr){};
